import json
from environment.trace_sample import Trace
from learning.eval_helper import Measurement
from config import DRLParameters
import numpy as np
import random
from typing import Tuple, List
from deprecated.sphinx import deprecated

random.seed(0)

data_set_name = "hospital_billing"  # TODO: 需修改
parameter_space = DRLParameters(dataset_name=data_set_name)


class Policy:
    """
    预测策略
    """

    def __init__(self):
        pass

    @staticmethod
    def predict_all_points(trace: Trace, alpha) -> Tuple[float, int]:
        """ 在每一点处都进行预测
        :param alpha:
        :param trace:
        :return: AI_loss, 预测的点个数
        """
        ytrue = trace.true_value
        all_pred = trace.prediction_sequence
        times = trace.time_sequence
        time_subseq = np.array(times)
        pred_subseq = np.array(all_pred)
        return Measurement.AILoss(time_subseq, pred_subseq, ytrue, alpha), len(time_subseq) - 1

    @staticmethod
    @deprecated(version="0.0.0", reason="此方法废弃，改为predict_random_points_v1")
    def predict_random_points(trace: Trace, max_step_length: int = 10) -> Tuple[float, int]:
        """ 在随机的位置进行预测，随机的最大步长为10
        :param trace:
        :param max_step_length: 最大的步长
        :return: AI_loss, 预测的点个数
        """
        ytrue = trace.true_value
        all_pred = trace.prediction_sequence
        times = trace.time_sequence
        index = 0
        time_subseq = [times[0]]
        pred_subseq = [all_pred[0]]
        while True:
            gap = random.randint(1, max_step_length)
            if index + gap >= trace.trace_length - 1:
                break
            time_subseq.append(times[index + gap])
            pred_subseq.append(all_pred[index + gap])
            index += gap
        time_subseq.append(times[-1])
        pred_subseq.append(all_pred[-1])
        time_subseq = np.array(time_subseq)
        pred_subseq = np.array(pred_subseq)
        return Measurement.AILoss(time_subseq, pred_subseq, ytrue), len(time_subseq) - 1

    @staticmethod
    def predict_random_points_v1(trace: Trace, alpha) -> Tuple[float, int]:
        """ 在每个位置随机选择是否进行预测，选择预测的概率默认为0.5
        :param alpha:
        :param trace:
        :return:
        """
        ytrue = trace.true_value
        all_pred = trace.prediction_sequence
        times = trace.time_sequence
        index = 0
        time_subseq = [times[0]]
        pred_subseq = [all_pred[0]]
        while True:
            if index >= trace.trace_length - 1:
                break
            if random.random() >= 0.5:
                time_subseq.append(times[index])
                pred_subseq.append(all_pred[index])
            index += 1
        time_subseq.append(times[-1])
        pred_subseq.append(all_pred[-1])
        time_subseq = np.array(time_subseq)
        pred_subseq = np.array(pred_subseq)
        return Measurement.AILoss(time_subseq, pred_subseq, ytrue, alpha), len(time_subseq) - 1

    @staticmethod
    def predict_each_few_points(trace: Trace, gap: int = 2, alpha=0) -> Tuple[float, int]:
        """ 每隔gap个点进行预测
        :param alpha:
        :param trace:
        :param gap: 预测的点位置间隔
        :return: AI_loss, 预测的点个数
        """
        ytrue = trace.true_value
        all_pred = trace.prediction_sequence
        times = trace.time_sequence
        index = 0
        time_subseq = [times[0]]
        pred_subseq = [all_pred[0]]
        while True:
            if index + gap >= trace.trace_length - 1:
                break
            time_subseq.append(times[index + gap])
            pred_subseq.append(all_pred[index + gap])
            index += gap
        time_subseq.append(times[-1])
        pred_subseq.append(all_pred[-1])
        time_subseq = np.array(time_subseq)
        pred_subseq = np.array(pred_subseq)
        return Measurement.AILoss(time_subseq, pred_subseq, ytrue, alpha), len(time_subseq) - 1

    @staticmethod
    def best_prediction_indexes(trace: Trace) -> List:
        """ 最佳预测点的下标
        :param trace:
        :return:
        """
        ytrue = trace.true_value
        all_pred = trace.prediction_sequence
        index = 1
        pred_indexes = [0]
        while index < trace.trace_length - 1:
            if abs(all_pred[index] - ytrue) < abs(all_pred[pred_indexes[-1]] - ytrue):
                pred_indexes.append(index)
            index += 1
        return pred_indexes

    @staticmethod
    def predict_lowest_uncertainty(trace: Trace, max_n_pres, alpha=0) -> Tuple[float, int]:
        act_error_path = parameter_space.ACTIVITY_ERRORS_PATH
        fr = open(act_error_path, "r")
        act_errors_dict = json.load(fr)
        sorted_dict = sorted(act_errors_dict.items(), key=lambda x: x[1], reverse=True)
        can_pre_acts = []
        for act_name, error in sorted_dict:
            can_pre_acts.append(act_name)
            if len(can_pre_acts) == max_n_pres:
                break

        ytrue = trace.true_value
        all_pred = trace.prediction_sequence
        times = trace.time_sequence
        index = 1
        time_subseq = [times[0]]
        pred_subseq = [all_pred[0]]
        while True:
            if index >= trace.trace_length - 1:
                break
            if trace.get_activity_names()[index] in can_pre_acts:
                time_subseq.append(times[index])
                pred_subseq.append(all_pred[index])
            index += 1
        time_subseq.append(times[-1])
        pred_subseq.append(all_pred[-1])
        time_subseq = np.array(time_subseq)
        pred_subseq = np.array(pred_subseq)
        return Measurement.AILoss(time_subseq, pred_subseq, ytrue, alpha), len(time_subseq) - 1


class ProcessBaseline:
    def __init__(self):
        self.losses = []
        self.n_preds = []

    def process(self, alpha):
        cids = Trace(parameter_space=parameter_space).test_case_ids

        cnt = 0
        for cid in cids:
            if parameter_space.DEBUG:
                if cid not in parameter_space.INITIAL_CASE_IDS:
                    continue
            cnt += 1
            trace = Trace(parameter_space=parameter_space, case_id=cid)
            # loss, n_pred = Policy.predict_all_points(trace, alpha=alpha)
            # loss, n_pred = Policy.predict_random_points_v1(trace, alpha=alpha)
            # loss, n_pred = Policy.predict_each_few_points(trace, gap=6, alpha=alpha)
            loss, n_pred = Policy.predict_lowest_uncertainty(trace, max_n_pres=5, alpha=alpha)  # TODO: 需修改
            self.losses.append(loss)
            self.n_preds.append(n_pred)

    def display_result(self):
        print("Length of results: ", len(self.losses))
        print("Avg loss: ", sum(self.losses) / len(self.losses),
              ". Avg no. points: ", sum(self.n_preds) / len(self.n_preds))


if __name__ == '__main__':
    alphas = [0, 0.2, 0.4, 0.6, 0.8, 1.0]
    for alpha in alphas:
        process_baseline = ProcessBaseline()
        process_baseline.process(alpha=alpha)
        process_baseline.display_result()
